import cv2
import numpy as np
import os
import pydicom as dicom
import pandas as pd
import shutil

from tqdm import tqdm


def dukeFileChoose(dir_path):
    '''
    即筛选出的脂肪饱和处理前序列文件名称列表、第二次增强后的文件名称列表
    :param dir_path: duke原数据的文件夹路径
    :return: dict{'pre':[],'post':[]}
    '''
    def judge_filename(filename):
        '''
        最终筛选出917个例子，3D形状从92 ~256

        用于下面特定的if条件，筛选出有ax 3d dyn且不包含['2nd','3rd','4th','1st']的文件名
        :param filename:
        :return:

        '''
        a=['2nd','3rd','4th','1st']
        for s in a:
            if s in filename:
                return True
        return False
    def judge_filename2(filename):
        '''
        用于下面特定的if条件，筛选出有ax 3d dyn且不包含['2nd','3rd','4th','1st']的文件名
        :param filename:
        :return:
        '''
        a=['']
        for s in a:
            if s in filename:
                return True
        return False
    global annotation_Box_path

    count = 0
    no_pre_count=0
    pre_path_list=[]
    post_path_list=[]
    case_len_list=[]#记录筛选出来的每个case的dcm文件数量
    for case in os.listdir(dir_path):
        if case=="Breast_MRI_151":#特殊处理，这个的pre序列和第二次增强对比序列的维度不一样
            continue
        b= os.listdir(os.path.join(dir_path, case))
        for tmp_dir in os.listdir(os.path.join(dir_path, case)):
            pre_sequence_dir = None
            second_enhanse_dir=None
            only_one=0
            file_list=os.listdir(os.path.join(dir_path, case, tmp_dir))
            okList=[]
            #-----找出对比前的影像-------#
            for sequence_dir in os.listdir(os.path.join(dir_path, case, tmp_dir)):
                # if  'pre' in sequence_dir \
                #         or '-ax 3d dyn' in sequence_dir or '-ax dynamic' in sequence_dir \
                #         or '-ax dyn' in sequence_dir:
                if '-Ax Vibrant' in sequence_dir \
                        or ('pre' in sequence_dir and 't1 pre' not in sequence_dir)\
                        or ('-ax 3d dyn' in sequence_dir and not judge_filename(sequence_dir))\
                        or '-ax dynamic' in sequence_dir \
                        or '-ax dyn-' in sequence_dir or\
                        '-ax 3d dyn pre' in sequence_dir:
                # if 'pre' in sequence_dir:
                # if '-ax 3d dyn' in sequence_dir:
                    pre_sequence_dir = sequence_dir
                    only_one+=1
                    okList.append(sequence_dir)
            if only_one!=1:
                if only_one==0:
                    no_pre_count+=1
                    print("没找到对比前序列", os.listdir(os.path.join(dir_path, case, tmp_dir)))
                else:
                    # print("有多个文件符合条件，无法确定哪个是对比前序列")
                    # print(case,okList)
                    pass
                continue
            #------------找出第二次增强对比后的影像-------------#
            only_one=0
            for sequence_dir in os.listdir(os.path.join(dir_path, case, tmp_dir)):
                if '2nd' in sequence_dir or 'Ph2Ax' in sequence_dir or '2ax' in sequence_dir:
                    only_one+=1
                    second_enhanse_dir=sequence_dir
            if only_one==1:#找到符合条件的对比前序列和第二次增强的对比后序列

                count += 1
            else:
                if only_one == 0:
                    print("没找到第二次增强序列")
                else:
                    print("有多个文件符合条件，无法确定哪个是第二次增强序列")
        #--处理筛选出来的数据--#
            post_path = os.path.join(dir_path, case, tmp_dir, second_enhanse_dir)
            pre_path = os.path.join(dir_path, case, tmp_dir, pre_sequence_dir)
            pre_path_list.append(pre_path)
            post_path_list.append(post_path)

            #对比组数据的dcm文件数量是否一致（也就是对比维度是否一致）
            # reader = sitk.ImageSeriesReader()
            # dicom_names = reader.GetGDCMSeriesFileNames(pre_path)
            # reader.SetFileNames(dicom_names)
            # image = reader.Execute()
            # image_array=sitk.GetArrayFromImage(image)
            # print(pre_path,image_array.shape)

    # print('序列数：', count)
    # print("筛选出来的数据的维度的范围",case_len_list[0],case_len_list[-1])
    return {"pre":pre_path_list,"post":post_path_list}
# 获取dicom文件的pixel_array
def get_pixel_array(dicom_file_path):
    dcm_ori = dicom.dcmread(dicom_file_path)
    arr = dcm_ori.pixel_array
    arr = arr.astype(np.int32)
    return arr


# 最大最小归一化
def pixel_array_normalize(pixel_array):
    return 255.0 * (pixel_array - pixel_array.min()) / (pixel_array.max() - pixel_array.min())


def get_subtraction(dicom1, dicom2):
    """
    获取剪影图像
    :param dicom1:  增强后的dicom文件路径
    :param dicom2:  增强前的dicom文件路径
    :return:        剪影图像矩阵
    """
    pa1 = get_pixel_array(dicom1)
    pa2 = get_pixel_array(dicom2)
    # pa1 = pixel_array_normalize(pa1)
    # pa2 = pixel_array_normalize(pa2)
    img = pa1 - pa2

    # 截断负值
    # img = np.clip(img, a_min=0, a_max=img.max())
    # img = pixel_array_normalize(img)
    # # 计算直方图
    # img = np.uint8(img)
    # hist = cv2.calcHist([img], [0], None, [255], [1, 256])
    # # plt.plot(hist)
    # # plt.show()
    # hist_sum = np.cumsum(hist)
    # # 截断归一化
    # min_val = 0
    # # min_val = np.searchsorted(hist_sum, hist_sum[-1] * 0.01, side='left')
    # max_val = np.searchsorted(hist_sum, hist_sum[-1] * 0.9998, side='left')
    # img = np.clip(img, a_min=min_val, a_max=max_val)
    # img = 255.0 * (img - min_val) / (max_val - min_val)

    return img


# 获取duke数据集的剪影图像
def subtractionDcmToArray(dir_path, save_path):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    L=dukeFileChoose(dir_path =dir_path)
    pre_file_list,post_file_list=L['pre'],L['post']
    count=0
    for i in tqdm(range(len(pre_file_list))):
        dcm_filename_list = os.listdir(pre_file_list[i])
        for dicom_file in dcm_filename_list:
            post_path = os.path.join(post_file_list[i], dicom_file)
            pre_path = os.path.join(pre_file_list[i], dicom_file)
            array_file_name=pre_path.split('/')[-4]+'-'+dicom_file[2:5]
            img = get_subtraction(post_path, pre_path)
            count+=1
            last_save_name=os.path.join(save_path,array_file_name+'.npy')
            np.save(last_save_name,img)
    print("total number:",count)




def get_center_slice(annotation_path):
    """
    获取每个病例病灶中心切片号
    :param annotation_path: 标注文件路径
    :return: 字典 key为病例ID value为切片号
    """
    annotations = pd.read_excel(annotation_path)
    caseId2sliceId = {}
    for annotation in annotations.itertuples():
        case_id = annotation[1]
        start_slice = annotation[6]
        end_slice = annotation[7]
        caseId2sliceId[case_id] = int((end_slice - start_slice) / 2 + start_slice)
    return caseId2sliceId


def draw_box(jpg_dir, img_with_box_save_dir,annotation_path):
    """
    画出标注框
    :param img_dir: 切片图片文件夹路径
    :param annotation_path: 标注文件路径
    :return: None
    """
    if not os.path.exists(img_with_box_save_dir):
        os.makedirs(img_with_box_save_dir)
    annotations = pd.read_excel(annotation_path)
    count=0

    for annotation in annotations.itertuples():
        case_id = annotation[1]
        start_row = annotation[2]
        start_column = annotation[4]
        end_row = annotation[3]
        end_column = annotation[5]
        start_slice = annotation[6]
        end_slice = annotation[7]
        sliceId = int((end_slice - start_slice) / 2 + start_slice)
        if sliceId < 100:
            sliceId = '0'+str(sliceId)
        else:
            sliceId = str(sliceId)
        read_name = os.path.join(jpg_dir, case_id + '_second_subtraction_norm_' + sliceId + '.jpg')
        if os.path.exists(read_name):
            print(read_name)
            img = cv2.imread(read_name)
            # img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

            cv2.rectangle(img, (start_column, start_row), (end_column, end_row), (0, 255, 0), 1)
            save_name=os.path.join(img_with_box_save_dir, os.path.basename(read_name))
            cv2.imwrite(save_name, img)
            count+=1
    print(count)


if __name__ == '__main__':

    # 转换duke数据集
    dir_path = '/data1/home/liukai/AllData/Duke/Duke/Duke-Breast-Cancer-MRI'
    save_path = '/data1/home/liukai/AllData/Duke/Duke/2dCutArray'
    subtractionDcmToArray(dir_path, save_path)
    # -------将jpg图片画上box标注并保存-----#

